from abc import ABC, abstractmethod
from typing import List, Tuple

# 模拟器抽象类
class Simulator(ABC):
    
    # 模拟计算量
    @abstractmethod
    def get_calculate_workload(self, tp_size: int) ->  List[float]: ...

    # 显存占用理论值 激活值*2 + 参数量 * 6
    @abstractmethod
    def get_communication_workload(self, tp_size: int) -> int: ...

    @abstractmethod
    def get_dp_communication_workload(self, tp_size: int, dp_size: int) -> float: ...

class CostCalculator:
    
    def __init__(self):
        pass

    # 统计计算量
    def statistic_single_layer_calculate_workload(self, tp_size: int) -> List[float]:
        # 分别表示 加、减、乘、除、指数 次数
        workload_total = [0, 0, 0, 0, 0]
        for attr_name, attr_value in self.__dict__.items():
            if attr_value is not None:
                workload_total = [x+y for x, y in zip(workload_total, attr_value.get_calculate_workload(tp_size))]
        
        return workload_total
    
    def statistic_single_layer_memory_workload(self, tp_size: int) -> Tuple[float, float]:
        # 返回该层的显存占用以及激活值占用
        return 0 ,0
    
    def statistic_single_layer_communication_workload(self, tp_size: int) -> float:
        communicate_workload = 0
        for attr_name, attr_value in self.__dict__.items():
            if attr_value is not None:
                communicate_workload += attr_value.get_communication_workload(tp_size)
        
        return communicate_workload
    
    def statistic_single_layer_dp_communication_workload(self, tp_size: int, dp_size: int) -> float:
        dp_communicate_workload = 0
        for attr_name, attr_value in self.__dict__.items():
            if attr_value is not None:
                dp_communicate_workload += attr_value.get_dp_communication_workload(tp_size, dp_size)
        
        return dp_communicate_workload


class MoeCostCalculator:
    def statistic_single_layer_calculate_workload_with_ep(self, tp_size: int, ep_size: int, experts_num: int) -> List[float]:
        return 0

    def statistic_single_layer_memory_workload_with_ep(self, tp_size: int, ep_size: int, experts_num: int) -> Tuple[float, float]:
        return 0, 0

    def statistic_single_layer_tp_communication_workload_with_ep(self, tp_size: int, ep_size: int, experts_num: int) -> float:
        communicate_workload = 0
        for attr_name, attr_value in self.__dict__.items():
            if attr_value is not None:
                communicate_workload += attr_value.get_communication_workload_with_ep(tp_size, ep_size, experts_num)
        
        return communicate_workload
    
    def statistic_single_layer_dp_communication_workload_with_ep(self, tp_size: int, dp_size: int, ep_size: int, experts_num: int) -> float:
        dp_communicate_workload = 0
        for attr_name, attr_value in self.__dict__.items():
            if attr_value is not None:
                dp_communicate_workload += attr_value.get_dp_communication_workload_with_ep(tp_size, dp_size, ep_size, experts_num)
        
        return dp_communicate_workload
